#include "stdstratego.h"
#include "stdsc.h"

#include "ilpstrategies.h"
#include "linebyline.h"

#define PREPROLOGUE 0
#define PROLOGUE    1
#define REPEAT      2
#define EPILOGUE    3

typedef struct skelcode_s
{
  code body;
  code split[];
  code single;
} skelcode_t;

stratego code metasplit(code body)
{
  strategies
    main =
        <table-put>("globals", "prologue", Compound([], []))
      ; <table-put>("globals", "repeat", Compound([], []))
      ; <table-put>("globals", "epilogue", Compound([], []))
      ; <topdown(try(MetaSplit <+ AddPromote))>(body)
      ; ![<id>, <table-get>("globals", "prologue")
              , <table-get>("globals", "repeat")
              , <table-get>("globals", "epilogue")
         ]
      ; map(\a -> AssignInit(CodeConst(a))\); \a -> ArrayInit(a)\; id

  rules
    MetaSplit:
      Stat(FunCall(Id(i), [CodeConst(c)])) -> Stat(EmptyExp)
      where ( <elem>(i, ["prologue", "repeat", "epilogue"])
              ; <table-put>("globals", i, c)
            )

    AddPromote:
      Declaration2(t@TypeSpec(a, b, c),i) ->
      Declaration2(DeclSpec([Promote], t), i)

    AddPromote:
      Declaration2(DeclSpec(d, t), i) ->
      Declaration2(DeclSpec([Promote|d], t), i)
}

static void findinitialruns(instance_t *instance, int instances)
{
  int ii, jj;

  for (ii=instances-1; ii >= 0; ii--)
  {
    instance[ii].runs = 0;

    /* Find number of initial runs */
    for (jj=0; jj < instance[ii].arguments; jj++)
    {
      if (instance[ii].argcode[jj].argtype == ARG_STREAM_OUT)
      {
        int ci = instance[ii].argcode[jj].connector.instance;
        int ca = instance[ii].argcode[jj].connector.argument;

        if (ci != -1)
        {
          instance[ii].runs = MAX(instance[ii].runs,
                                  instance[ci].argcode[ca].tokens.request -
                                  instance[ii].argcode[jj].tokens.initial);
        }
      }
    }

    /* Find how much tokens we produce on our outputs, and
     * how much data we need on our inputs */
    for (jj=0; jj < instance[ii].arguments; jj++)
    {
      if (instance[ii].argcode[jj].argtype == ARG_STREAM_IN ||
          instance[ii].argcode[jj].argtype == ARG_STREAM_OUT)
      {
        instance[ii].argcode[jj].tokens.request =
          instance[ii].argcode[jj].tokens.initial + instance[ii].runs;
      }
    }
  }

}

static void findfinalruns(instance_t *instance, int instances)     
{
  int ii, jj;

  for (ii=0; ii < instances; ii++)
  {
    instance[ii].runs = 0;

    /* Find number of final runs */
    for (jj=0; jj < instance[ii].arguments; jj++)
    {
      if (instance[ii].argcode[jj].argtype == ARG_STREAM_IN)
      {
        int ci = instance[ii].argcode[jj].connector.instance;
        int ca = instance[ii].argcode[jj].connector.argument;

        if (ci != -1)
        {
          instance[ii].runs = MAX(instance[ii].runs,
                                  instance[ci].argcode[ca].tokens.request -
                                  instance[ii].argcode[jj].tokens.final);
        }
      }
    }

    /* Find how much tokens we produce on our outputs, and
     * how much data we need on our inputs */
    for (jj=0; jj < instance[ii].arguments; jj++)
    {
      if (instance[ii].argcode[jj].argtype == ARG_STREAM_IN ||
          instance[ii].argcode[jj].argtype == ARG_STREAM_OUT)
      {
        instance[ii].argcode[jj].tokens.request =
          instance[ii].argcode[jj].tokens.initial + instance[ii].runs;
      }
    }
  }
}

static void updatebufferfromruns(instance_t *instance, int instances)
{
  int ii, jj;

  /* Update buffer sizes */
  for (ii=0; ii < instances; ii++)
  {
    for (jj=0; jj < instance[ii].arguments; jj++)
    {
      if (instance[ii].argcode[jj].argtype == ARG_STREAM_IN ||
          instance[ii].argcode[jj].argtype == ARG_STREAM_OUT)
      {
        int ci = instance[ii].argcode[jj].connector.instance;
        int ca = instance[ii].argcode[jj].connector.argument;

        if (ci == -1)
        {
          instance[ii].argcode[jj].buffer.size =
            MAX(instance[ii].argcode[jj].buffer.size,
                instance[ii].argcode[jj].tokens.request);
        }
        else
        {
          instance[ii].argcode[jj].buffer.size =
            MAX(instance[ii].argcode[jj].buffer.size,
                MAX(instance[ii].argcode[jj].tokens.request,
                    instance[ci].argcode[ca].tokens.request));
        }
      }
    }
  }
}

static void findstate(instance_t *instance, int instances)
{
  int ii, jj;

  /* Always need at least one line, for repetitive part */
  for (ii=0; ii < instances; ii++)
    for (jj=0; jj < instance[ii].arguments; jj++)
      instance[ii].argcode[jj].buffer.size = 1;

  findinitialruns(instance, instances);
  updatebufferfromruns(instance, instances);
  findfinalruns(instance, instances);
  updatebufferfromruns(instance, instances);

  for (ii=0; ii < instances; ii++)
    for (jj=0; jj < instance[ii].arguments; jj++)
    {
      int ci = instance[ii].argcode[jj].connector.instance;
      int ca = instance[ii].argcode[jj].connector.argument;

      if (ci != -1)
      {
        instance[ii].argcode[jj].buffer.size =
          MAX(instance[ii].argcode[jj].buffer.size,
              instance[ci].argcode[ca].buffer.size);

        if (instance[ii].argcode[jj].argtype == ARG_STREAM_OUT)
        {
          instance[ii].argcode[jj].buffer.data = makeid(makeid(makeid(`__buffer_i`, ii), `_a`), jj);
          instance[ii].argcode[jj].buffer.rp = makeid(instance[ii].argcode[jj].buffer.data, `_rp`);
          instance[ii].argcode[jj].buffer.wp = makeid(instance[ii].argcode[jj].buffer.data, `_wp`);
        }

        if (instance[ii].argcode[jj].argtype == ARG_STREAM_IN)
        {
          instance[ii].argcode[jj].buffer.data = instance[ci].argcode[ca].buffer.data;
          instance[ii].argcode[jj].buffer.rp = instance[ci].argcode[ca].buffer.rp;
          instance[ii].argcode[jj].buffer.wp = instance[ci].argcode[ca].buffer.wp;
        }
      }
    }

  /* Lose one line due to empty/full ambiguity */
  for (ii=0; ii < instances; ii++)
    for (jj=0; jj < instance[ii].arguments; jj++)
       instance[ii].argcode[jj].buffer.size += 1;
}

static void declarebuffers(instance_t *instance, int instances)
{
  int ii, jj, kk;

  for (ii=0; ii < instances; ii++)
    for (jj=0; jj < instance[ii].arguments; jj++)
    {
      int ci = instance[ii].argcode[jj].connector.instance;
      int ca = instance[ii].argcode[jj].connector.argument;

      if (ci != -1)
      {
        if (instance[ii].argcode[jj].argtype == ARG_STREAM_OUT)
        {
          promote int @instance[ii].argcode[jj].buffer.rp = 0;
          promote int @instance[ii].argcode[jj].buffer.wp = 0;

          @declare1d(addpointer(instance[ii].argcode[jj].datatype),
                     instance[ii].argcode[jj].buffer.data, instance[ii].argcode[jj].buffer.size);
          
          for (kk=0; kk < instance[ii].argcode[jj].buffer.size; kk++)
          {
            (@instance[ii].argcode[jj].buffer.data)[kk] =
              @cast(addpointer(instance[ii].argcode[jj].datatype),
                    malloc(stride*sizeof(@instance[ii].argcode[jj].datatype)));
          }
 
        }
      }
    }  
}

/* For all buffer interactions:
 * Dynamic scope: args, instance, instances, skelcode
 */

static int instanceBufferGetStride(int ii, int aa)
{
  if (instance[ii].argcode[aa].connector.instance == -1)
    return scuBufferGetStride(args[instance[ii].argcode[aa].connector.argument].buffer);
  else
    /* Assumes all strides are equal */
    return scuBufferGetStride(args[instance[fii].argcode[fia].connector.argument].buffer);
}

static int instanceBufferGetLines(int ii, int aa)
{
  if (instance[ii].argcode[aa].connector.instance == -1)
    return scuBufferGetLines(args[instance[ii].argcode[aa].connector.argument].buffer);
  else
    /* Assumes all heights are equal */
    return scuBufferGetLines(args[instance[fii].argcode[fia].connector.argument].buffer);
}

static int instanceBufferRemoveReference(int ii, int aa)
{
  if (instance[ii].argcode[aa].connector.instance == -1)
    return scuBufferRemoveReference(args[instance[ii].argcode[aa].connector.argument].buffer);
  else
    return 0;
}

static int instanceBufferFinalize(int ii, int aa)
{
  if (instance[ii].argcode[aa].connector.instance == -1)
    return scuBufferFinalize(args[instance[ii].argcode[aa].connector.argument].buffer);
  else
    return 0;
}

static int instanceBufferPeek(int ii, int aa, void **data, int size, int options)
{
  if (instance[ii].argcode[aa].connector.instance == -1)
    return scuBufferPeek(args[instance[ii].argcode[aa].connector.argument].buffer,
                         args[instance[ii].argcode[aa].connector.argument].reader,
                        data, size, options);
  else
  {
    *data = (@instance[ii].argcode[aa].buffer.data)[@instance[ii].argcode[aa].buffer.rp];
    (@instance[ii].argcode[aa].buffer.rp)++;
    if (@instance[ii].argcode[aa].buffer.rp == instance[ii].argcode[aa].buffer.size)
      @instance[ii].argcode[aa].buffer.rp = 0;

    return size;
  }
}

static void instanceBufferReleasePeeked(int ii, int aa, int size)
{
  if (instance[ii].argcode[aa].connector.instance == -1)
    scuBufferReleasePeeked(args[instance[ii].argcode[aa].connector.argument].buffer,
                           args[instance[ii].argcode[aa].connector.argument].reader,
                           size);
}

static int instanceBufferAllocate(int ii, int aa, void **data, int size, int options)
{
  if (instance[ii].argcode[aa].connector.instance == -1)
    return scuBufferAllocate(args[instance[ii].argcode[aa].connector.argument].buffer,
                             data, size, options);
  else
  {
    *data = (@instance[ii].argcode[aa].buffer.data)[@instance[ii].argcode[aa].buffer.wp];
    (@instance[ii].argcode[aa].buffer.wp)++;
    if (@instance[ii].argcode[aa].buffer.wp == instance[ii].argcode[aa].buffer.size)
      @instance[ii].argcode[aa].buffer.wp = 0;

    return size;
  }
}

static void instanceBufferReleaseAllocated(int ii, int aa, int size)
{
  if (instance[ii].argcode[aa].connector.instance == -1)
    scuBufferReleaseAllocated(args[instance[ii].argcode[aa].connector.argument].buffer,
                              size);
}

static code rewriteBuffers(int ii, code body)
{
  body = rewrite(`BufferGetStride(PH1)`, `instanceBufferGetStride($ii, PH1)`, body);
  body = rewrite(`BufferGetLines(PH1)`, `instanceBufferGetLines($ii, PH1)`, body);
  body = rewrite(`BufferRemoveReference(PH1)`, `instanceBufferRemoveReference($ii, PH1)`, body);
  body = rewrite(`BufferFinalize(PH1)`, `instanceBufferFinalize($ii, PH1)`, body);
  body = rewrite(`BufferPeek(PH1, PH2, PH3, PH4)`, `instanceBufferPeek($ii, PH1, PH2, PH3, PH4)`, body);
  body = rewrite(`BufferReleasePeeked(PH1, PH2)`, `instanceBufferReleasePeeked($ii, PH1, PH2)`, body);
  body = rewrite(`BufferAllocate(PH1, PH2, PH3, PH4)`, `instanceBufferAllocate($ii, PH1, PH2, PH3, PH4)`, body);
  body = rewrite(`BufferReleaseAllocated(PH1, PH2)`, `instanceBufferReleaseAllocated($ii, PH1, PH2)`, body);

  return body;
}

void linebyline(instance_t *instance, int instances)
{
  skelcode_t skelcode[];
  int ii, jj, fii=-1, fia=-1;
  int lbl_hpixels, stride;

  for (ii=0; ii < instances; ii++)
  {
    /* Make sure function calls in registered codes are expanded */
    instance[ii].body = rewrite(`prologue(PH1)`, `prologue(reduce(PH1, 1))`, instance[ii].body);
    instance[ii].body = rewrite(`repeat(PH1)`, `repeat(reduceit(PH1, 1))`, instance[ii].body);
    instance[ii].body = rewrite(`epilogue(PH1)`, `epilogue(reduce(PH1, 1))`, instance[ii].body);

    /* Add skeleton functions to symbol table */
    @addpromote(instance[ii].body);

    /* Call skeleton, keeping expanded code */
    skelcode[ii].body = reduce(`($@instance[ii].skeleton)(instance[ii].argcode, instance[ii].arguments, instance[ii].kernel)`);

    /* Extract registered codes, and make variables unique */
    skelcode[ii].split = @metasplit(skelcode[ii].body);
    skelcode[ii].split = @scramblevariables(skelcode[ii].split, makeid(`i`, ii));

    /* Add instance information to buffer interactions */
    for (jj=0; jj < 4; jj++)
      if (jj != REPEAT)
        skelcode[ii].split[jj] = rewriteBuffers(ii, skelcode[ii].split[jj]);

    /* Rename arguments to match scrambled variables */
    for (jj=0; jj < instance[ii].arguments; jj++)
    {
      instance[ii].argcode[jj].id = makeid(makeid(makeid(`i`, ii), instance[ii].argcode[jj].id), `0`);

      if (instance[ii].argcode[jj].connector.instance == -1 &&
          instance[ii].argcode[jj].argtype == 0 && fia == -1)
      {
        fii = ii;
        fia = jj;
      }
    }
  }
  
  stride = instanceBufferGetStride(fii, fia);
  lbl_hpixels = stride;

  /* Preprologue (initialize variables, no buffer interaction) */
  for (ii=0; ii < instances; ii++)
    @skelcode[ii].split[0];

  findstate(instance, instances);
  declarebuffers(instance, instances);

  /* Prepare code for processing a single line */
  if (instances > 1)
  {
    for (ii=0; ii < instances; ii++)
    {
      skelcode[ii].single = reduceit(`singleline(instance[$ii].argcode, instance[$ii].arguments, skelcode[$ii].split[REPEAT]);`, 0);
      skelcode[ii].single = rewriteBuffers(ii, skelcode[ii].single);
    }
  }

  /* Set properties */
  propEnvReset();
  propEnv("operation", operation);
 
  for (ii=0; ii < instances; ii++)
    for (jj=0; jj < instance[ii].arguments; jj++)
      if (instance[ii].argcode[jj].connector.instance == -1)
      {
        int ca = instance[ii].argcode[jj].connector.argument;

        propEnv("argument", @idtostring(makeid(`id`, ca)));
          propSeti("id", ca);
          propEnv("target", target);
            propSeti("buffer", instance[ii].argcode[jj].state.buffer);
            propSeti("state", instance[ii].argcode[jj].state.state);
            propSeti("delay", instance[ii].argcode[jj].state.delay);
            propSeti("elementsize", sizeof(instance[ii].argcode[jj].datatype));

            if (instance[ii].argcode[jj].argtype == ARG_STREAM_IN ||
                instance[ii].argcode[jj].argtype == ARG_STREAM_OUT ||
                instance[ii].argcode[jj].argtype == ARG_LOOKUP)
              propSeti("chunksize", 1024);
            else
              propSeti("chunksize", 1);

            if (instance[ii].argcode[jj].argtype == ARG_SCALAR_IN ||
	        instance[ii].argcode[jj].argtype == ARG_SCALAR_OUT || 
                instance[ii].argcode[jj].argtype == ARG_LOOKUP)
              propSets("timedist", "bulk");
            else
              propSets("timedist", "regular");
          propEnvBack();
        propEnvBack();
      }

  /* Prologue (Preallocate buffers) 
   * Body must be able to run without prologue */
  findinitialruns(instance, instances);

  for (ii=0; ii < instances; ii++)
  {
    @skelcode[ii].split[1];
    for (jj=0; jj < instance[ii].runs; jj++)
      @skelcode[ii].single;
  }

  /* Actual processing */
  while (lbl_hpixels = instanceBufferPeek(fii, fia,
                       (void**)&(@instance[fii].argcode[fia].id)[0], stride*sizeof(instance[fii].argcode[fia].datatype), 1))
  {
    lbl_hpixels = lbl_hpixels / sizeof(instance[fii].argcode[fia].datatype);

    for (ii=0; ii < instances; ii++)
      for (jj=0; jj < instance[ii].arguments; jj++)
        if ((ii != fii || jj != fia) && instance[ii].argcode[jj].argtype == 0)
          instanceBufferPeek(ii, jj, (void**)&(@instance[ii].argcode[jj].id)[0],
                     lbl_hpixels*sizeof(instance[ii].argcode[jj].datatype), 1);
        else if (instance[ii].argcode[jj].argtype == 1)
          instanceBufferAllocate(ii, jj, (void**)&(@instance[ii].argcode[jj].id)[0], lbl_hpixels*sizeof(instance[ii].argcode[jj].datatype),
                         1);

    for (ii=0; ii < instances; ii++)
      @skelcode[ii].split[2];

    for (ii=0; ii < instances; ii++)
      for (jj=0; jj < instance[ii].arguments; jj++)
        if (instance[ii].argcode[jj].argtype == 0)
          instanceBufferReleasePeeked(ii, jj, lbl_hpixels*sizeof(instance[ii].argcode[jj].datatype));
        else if (instance[ii].argcode[jj].argtype == 1)
          instanceBufferReleaseAllocated(ii, jj, lbl_hpixels*sizeof(instance[ii].argcode[jj].datatype));
  }

  lbl_hpixels = stride;

  /* Epilogue (release buffers) */
  findfinalruns(instance, instances);

  for (ii=0; ii < instances; ii++)
  {
    for (jj=0; jj < instance[ii].runs; jj++)
      @skelcode[ii].single;

    @skelcode[ii].split[3];
  }

}

static void singleline(argcode_t *argcode, int arguments, code body)
{
  int jj;

  for (jj=0; jj < arguments; jj++)
    if (argcode[jj].argtype == ARG_STREAM_IN)
      BufferPeek(jj, (void**)&(@argcode[jj].id)[0],
                 lbl_hpixels*sizeof(argcode[jj].datatype), 1);
    else if (argcode[jj].argtype == ARG_STREAM_OUT)
      BufferAllocate(jj, (void**)&(@argcode[jj].id)[0],
                     lbl_hpixels*sizeof(argcode[jj].datatype), 1);

  @body;

  for (jj=0; jj < arguments; jj++)
    if (argcode[jj].argtype == ARG_STREAM_IN)
      BufferReleasePeeked(jj, lbl_hpixels*sizeof(argcode[jj].datatype));
    else if (argcode[jj].argtype == ARG_STREAM_OUT)
      BufferReleaseAllocated(jj, lbl_hpixels*sizeof(argcode[jj].datatype));
}
